package com.lingjtx.common.thread.mq;

import com.lingjtx.common.thread.pooling.ObjectPoolTemplate;
import com.rabbitmq.stream.Environment;
import com.rabbitmq.stream.Producer;
import jakarta.annotation.PreDestroy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessageProperties;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

public class RabbitStreamProducer {
    private static final Logger log = LoggerFactory.getLogger(RabbitStreamProducer.class);

    private final ObjectPoolTemplate<Environment> environmentObjectPoolTemplate;
    private final FastJsonAutoTypeMessageConverter converter;

    // 缓存 producer，key 是 stream 名称
    private final Map<String, Producer> producerCache = new ConcurrentHashMap<>();


    public RabbitStreamProducer(ObjectPoolTemplate<Environment> environmentObjectPoolTemplate, FastJsonAutoTypeMessageConverter converter) {
        this.environmentObjectPoolTemplate = environmentObjectPoolTemplate;
        this.converter = converter;
    }

    private Producer getOrCreateProducer(String streamName) {
        return producerCache.computeIfAbsent(streamName, name ->
                environmentObjectPoolTemplate.execute(environment -> environment.producerBuilder()
                        .stream(name)
                        .build())
        );
    }

    public void send(String streamName, Object message) {
        try (Producer producer = getOrCreateProducer(streamName)) {
            MessageProperties props = new MessageProperties();
            Message amqpMessage = converter.toMessage(message, props);
            producer.send(producer.messageBuilder()
                    .addData(amqpMessage.getBody())
                    .build(), confirmationStatus -> {
                if (!confirmationStatus.isConfirmed()) {
                    log.error("Message failed to reach stream: {}", streamName);
                }
            });
        } catch (Exception e) {
            throw new RuntimeException("消息发送失败", e);
        }
    }

    public void asyncSend(String streamName, Object message, Runnable onSuccess, Consumer<Throwable> onError) {
        try (Producer producer = getOrCreateProducer(streamName)) {
            MessageProperties props = new MessageProperties();
            Message amqpMessage = converter.toMessage(message, props);
            producer.send(producer.messageBuilder()
                    .addData(amqpMessage.getBody())
                    .build(), status -> {
                if (status.isConfirmed()) {
                    onSuccess.run();
                } else {
                    onError.accept(new RuntimeException("发送未确认"));
                }
            });
        } catch (Exception e) {
            onError.accept(e);
        }
    }

    @PreDestroy
    public void shutdown() {
        producerCache.values().forEach(Producer::close);
        producerCache.clear();
        log.info("🛑 所有 RabbitMQ Stream 生产者已关闭");
    }
}
